package com.shi.aspect;

import com.shi.annotation.limit.RateLimitRule;
import com.shi.annotation.limit.RateLimiter;
import com.shi.annotation.limit.RateLimiters;
import com.shi.enums.LimitType;
import com.shi.exception.BusinessException;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Component
@Aspect
@Order(2)
public class RateLimiterAspect {
    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private DefaultRedisScript<Long> limitScript;

    {
        limitScript = new DefaultRedisScript<>("""
                local flag = 1
                for i = 1, #KEYS do
                    local window_start = tonumber(ARGV[1]) - tonumber(ARGV[(i - 1) * 2 + 2])
                    print(tonumber(ARGV[(i - 1) * 2 + 2]))
                    redis.call('ZREMRANGEBYSCORE', KEYS[i], 0, window_start)
                    local current_requests = tonumber(redis.call('ZCARD', KEYS[i]))
                    if current_requests < tonumber(ARGV[(i - 1) * 2 + 3]) then
                    else
                        flag = 0
                    end
                end
                if flag == 1 then
                    for i = 1, #KEYS do
                        print('add')
                        print(KEYS[i], tonumber(ARGV[1]), ARGV[1])
                        redis.call('ZADD', KEYS[i], tonumber(ARGV[1]), ARGV[1])
                        redis.call('pexpire', KEYS[i], tonumber(ARGV[(i - 1) * 2 + 2]))
                    end
                end
                return flag
                """, Long.class);
    }

    private static String getIpAddr(HttpServletRequest request) {
        String ipAddress;
        try {
            ipAddress = request.getHeader("x-forwarded-for");
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getRemoteAddr();
                if (ipAddress.equals("127.0.0.1")) {
                    // 根据网卡取本机配置的IP
                    try {
                        ipAddress = InetAddress.getLocalHost().getHostAddress();
                    } catch (BusinessException e) {
                        throw new RuntimeException();
                    }
                }
            }
            // 通过多个代理的情况，第一个IP为客户端真实IP,多个IP按照','分割
            if (ipAddress != null) {
                if (ipAddress.contains(",")) {
                    return ipAddress.split(",")[0];
                } else {
                    return ipAddress;
                }
            } else {
                return "";
            }
        } catch (Exception e) {
            e.printStackTrace();
            return "";
        }
    }

    // 定义切点，需要把RateLimiter和RateLimiters同时加进来，否则多重注解不生效
    @Pointcut("@annotation(com.shi.annotation.limit.RateLimiter)")
    public void rateLimiter() {
    }

    @Pointcut("@annotation(com.shi.annotation.limit.RateLimiters)")
    public void rateLimiters() {
    }

    // 定义切点之前的操作
    @Before("rateLimiter() || rateLimiters()")
    public void doBefore(JoinPoint point) {
        try {
            // 从切点获取方法签名
            MethodSignature signature = (MethodSignature) point.getSignature();
            // 获取方法
            Method method = signature.getMethod();
            String name = point.getTarget().getClass().getName() + "." + signature.getName();
            // 获取日志注解
            RateLimiter rateLimiter = method.getAnnotation(RateLimiter.class);
            RateLimiters rateLimiters = method.getAnnotation(RateLimiters.class);

            List<RateLimiter> limiters = new ArrayList<>();
            if (rateLimiter != null) {
                limiters.add(rateLimiter);
            }

            if (rateLimiters != null) {
                limiters.addAll(Arrays.asList(rateLimiters.value()));
            }

            if (!allowRequest(limiters, name)) {
                throw new RuntimeException("访问过于频繁，请稍候再试");
            }

        } catch (RuntimeException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException("服务器限流异常，请稍候再试");
        }
    }

    /**
     * 是否允许请求
     *
     * @param rateLimiters 限流注解
     * @param name         方法全名
     * @return 是否放行
     */
    private boolean allowRequest(List<RateLimiter> rateLimiters, String name) {
        List<String> keys = getKeys(rateLimiters, name);
        String[] args = getArgs(rateLimiters);
        Long res = stringRedisTemplate.execute(limitScript, keys, args);
        System.out.println(res);
        return res != null && res == 1L;
    }

    /**
     * 获取限流的键
     *
     * @param rateLimiters 限流注解
     * @param name         方法全名
     * @return
     */
    private List<String> getKeys(List<RateLimiter> rateLimiters, String name) {
        List<String> keys = new ArrayList<>();

        for (RateLimiter rateLimiter : rateLimiters) {
            String key = rateLimiter.key();
            RateLimitRule[] rules = rateLimiter.rules();
            LimitType type = rateLimiter.type();

            StringBuilder sb = new StringBuilder();
            sb.append(key).append(name);

            if (LimitType.IP == type) {
                String ipAddr = getIpAddr(((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest());
                sb.append("_").append(ipAddr);
            } else if (LimitType.USER == type) {
                // 获取用户id,自己实现
                Long userId = 1L;
                sb.append("_").append(userId);
            }
            for (RateLimitRule rule : rules) {
                int time = rule.time() * 1000;
                int count = rule.count();
                StringBuilder builder = new StringBuilder(sb);
                builder.append("_").append(time).append("_").append(count);
                keys.add(builder.toString());
            }
        }
        return keys;
    }

    /**
     * 获取需要的参数
     *
     * @param rateLimiters 限流注解
     * @return
     */
    private String[] getArgs(List<RateLimiter> rateLimiters) {
        List<String> args = new ArrayList<>();
        args.add(String.valueOf(System.currentTimeMillis()));
        for (RateLimiter rateLimiter : rateLimiters) {
            RateLimitRule[] rules = rateLimiter.rules();
            for (RateLimitRule rule : rules) {
                int time = rule.time() * 1000;
                int count = rule.count();
                args.add(String.valueOf(time));
                args.add(String.valueOf(count));
            }
        }
        return args.toArray(new String[]{});
    }


}